# generate predicted scores given fitted idealstan model
# then export to Stata

library(tidyverse)
library(idealstan)
library(haven)
library(posterior)

# please set these locations first to replication files
# load replicated idealstan fit

save_loc <- "~/wbsurvey/Replication Package 2/In"

wb_fit <- readRDS(paste0(save_loc,"/wb_fit_replication.rds"))

# extract covariates

wb_cov <- select(wb_fit@score_data@score_matrix,`agg_sectorMedium-high-technology Manufacturing`:v2x_polyarchy) %>% 
  as.matrix

# extract group_level intercepts

wb_group <- wb_fit@stan_samples$draws("L_full") %>% 
  as_draws_matrix

# extract covariate coefficients

wb_cov_coef <- wb_fit@stan_samples$draws("legis_x") %>% 
  as_draws_matrix

# iterate over rows, use country-level intercept and add 
# covariate influence with matrix multiplicate
# for each draw

over_rows <- lapply(1:nrow(wb_fit@score_data@score_matrix), function(i) {
  
  this_row <- wb_group[,as.numeric(wb_fit@score_data@score_matrix$group_id[i])] + t(wb_cov[i,,drop=FALSE] %*% t(wb_cov_coef))
  
  tibble(person_id=wb_fit@score_data@score_matrix$person_id[i],
         group_id=wb_fit@score_data@score_matrix$group_id[i],
         score_raw=as.numeric(this_row),
         score_trans=plogis(as.numeric(this_row))) %>% 
    mutate(draw=as.numeric(attr(this_row, "dimnames")$draw))
  
}) %>% bind_rows

# save raw file

over_rows %>% 
  select(idstd="person_id",
         country="group_id",
         score_raw,score_trans) %>% 
  saveRDS(paste0(save_loc,"/Processed data/scores_raw.rds"))

# update Stata dataset with replicated scores

stata_dta <- read_dta(paste0(save_loc,"/Processed data/data_for_analysis.dta"))

to_stata <- over_rows %>% 
  select(idstd="person_id",
         score_trans) %>% 
  mutate(idstd=as.numeric(as.character(idstd))) %>% 
  group_by(idstd) %>% 
  summarize(firm_score_est=mean(score_trans*100),
            firm_score_95=quantile(score_trans*100,.95),
            firm_score_05=quantile(score_trans*100,.05),
            firm_score_se=sd(score_trans*100),
            firm_score_var=var(score_trans*100))

# save scores

saveRDS(to_stata,paste0(save_loc,"/Processed data/scores_sum.rds"))
write.csv(to_stata,paste0(save_loc,"/Processed data/scores_sum.csv"))

stata_dta <- select(stata_dta, 
                     -one_of(c("firm_score_est",
                             "firm_score_95",
                             "firm_score_05",
                             "firm_score_se",
                             "firm_score_var")))

stata_dta <- left_join(stata_dta, 
                       to_stata,
                       by="idstd")

write_dta(stata_dta, paste0(save_loc,"/Processed data/data_for_analysis.dta"))
